// Connect.cpp : Implementation of CConnect
#include "stdafx.h"
#include "AddIn.h"
#include "TestSuiteTool.h"
#include "Connect.h"

extern CAddInModule _AtlModule;

const wchar_t * const TOOLBAR_NAME = L"CppClean";
const wchar_t * const OUTPANE_NAME = L"CppClean";
const wchar_t * const MODULE_NAME = L"CppClean.Connect.";

static CConnect::CommandInfo s_commandList[] =
{
	{L"CreateTestProject", L"", true,  false, IDB_CREATE_TESTPROJECT, IDS_CREATE_TESTPROJECT},
	{L"RunTestProject", L"", true,  false, IDB_RUN_TESTPROJECT, IDS_RUN_TESTPROJECT},
	{L"AutoMock", L"", true,  false, IDB_AUTO_MOCK, IDS_AUTO_MOCK},
};

//////////////////////////////////////////////////////////////////////////
//
//////////////////////////////////////////////////////////////////////////
CConnect::CConnect()
{
	m_pTestSuite = new CTestSuiteTool(this);
}

CConnect::~CConnect()
{
	if(m_pTestSuite)
	{
		delete m_pTestSuite;
		m_pTestSuite = NULL;
	}
}

int CConnect::myMessageBox(UINT uID, UINT uType)
{
	HINSTANCE hInst = _AtlModule.GetResourceInstance();
	TCHAR szMessage[MAX_STRING_LEN] = {0};
	LoadString(hInst, uID, szMessage, MAX_STRING_LEN);
	TCHAR szTitle[MAX_STRING_LEN] = {0};
	LoadString(hInst, IDS_PROJNAME, szTitle, MAX_STRING_LEN);

	return ::MessageBox(GetActiveWindow(), szMessage, szTitle, uType);
}

int CConnect::myMessageBox(LPCTSTR str, UINT uType)
{
	HINSTANCE hInst = _AtlModule.GetResourceInstance();
	TCHAR szTitle[MAX_STRING_LEN] = {0};
	LoadString(hInst, IDS_PROJNAME, szTitle, MAX_STRING_LEN);

	return ::MessageBox(GetActiveWindow(), str, szTitle, uType);
}

wstring CConnect::myLoadString(UINT uID)
{
	HINSTANCE hInst = _AtlModule.GetResourceInstance();
	WCHAR szMessage[MAX_STRING_LEN] = {0};
	LoadString(hInst, uID, szMessage, MAX_STRING_LEN);
	return (szMessage);
}

BOOL CConnect::myExtractSubString(wstring& rString, LPCTSTR lpszFullString, int iSubString, TCHAR chSep)
{
	if (lpszFullString == NULL)
		return FALSE;

	while (iSubString--)
	{
		lpszFullString = _tcschr(lpszFullString, chSep);
		if (lpszFullString == NULL)
		{
			rString.clear();        // return empty string as well
			return FALSE;
		}
		lpszFullString++;       // point past the separator
	}
	LPCTSTR lpchEnd = _tcschr(lpszFullString, chSep);
	int nLen = (lpchEnd == NULL) ? lstrlen(lpszFullString) : (int)(lpchEnd - lpszFullString);

	TCHAR* pTemp = new TCHAR[nLen + 1];
	memset(pTemp, 0, (nLen + 1) * sizeof(TCHAR));
	memcpy(pTemp, lpszFullString, nLen*sizeof(TCHAR));
	rString = pTemp;

	return TRUE;
}

// When run, the Add-in wizard prepared the registry for the Add-in.
// At a later time, if the Add-in becomes unavailable for reasons such as:
//   1) You moved this project to a computer other than which is was originally created on.
//   2) You chose 'Yes' when presented with a message asking if you wish to remove the Add-in.
//   3) Registry corruption.
// you will need to re-register the Add-in by building the MyAddin21Setup project 
// by right clicking the project in the Solution Explorer, then choosing install.
void UnregisterCommand(CComPtr<EnvDTE::Commands>& pCommands, LPCWSTR commandName)
{
	CComPtr<EnvDTE::Command> pCommand;

	//
	// COMMAND_NAME_FULL must be the module name plus the COMMAND_NAME.
	// For this case, the module name can be found at the CvsInVC7.rgs and
	// CvsInVC8.rgs files.
	//
	WCHAR item[256];
	wcscpy_s(item, 256, MODULE_NAME);
	wcscat_s(item, 256, commandName);

	HRESULT hr = pCommands->Item(CComVariant(item), -1, &pCommand);

	if (SUCCEEDED(hr))
	{
		pCommand->Delete();
	}
}

void UnregisterAllCommands(CComPtr<DTE2>& pDTE)
{
	CComPtr<EnvDTE::Commands> pCommands;
	pDTE->get_Commands(&pCommands);

	int nMaxCommand = sizeof(s_commandList)/sizeof(s_commandList[0]);
	for(int curCommand = 0;	curCommand < nMaxCommand; ++curCommand)
	{
		const CConnect::CommandInfo * commandInfo = &s_commandList[curCommand];
		UnregisterCommand(pCommands, commandInfo->m_name);
	}
}

void UnregisterAllCommandBars(CComPtr<DTE2>& pDTE)
{
	CComPtr<IDispatch> ppv;

	// Get the set of command bars for the application.
	HRESULT hr = pDTE->get_CommandBars(&ppv);
	if (SUCCEEDED(hr))
	{
		CComQIPtr<_CommandBars> pCommandBars(ppv);
		if (pCommandBars)
		{
			CComPtr<CommandBar> pCommandBar;
			// See if the CvsIn toolbar has been created.
			hr = CConnect::FindToolbarByName(pCommandBars, TOOLBAR_NAME, &pCommandBar);
			if (SUCCEEDED(hr))
			{
				pCommandBar->Delete();
			}
		}
	}
}

void ConnectAndUnregisterAllCommands()
{
	// Create the shell and get an interface to it.
	CComPtr<EnvDTE::_DTE> pDTE;
	CComPtr<DTE2> pDTE2;
	HRESULT hr = pDTE.CoCreateInstance(L"VisualStudio.DTE", 0, CLSCTX_ALL);
	hr = pDTE.QueryInterface(&pDTE2);

	if(SUCCEEDED(hr) && pDTE2)
	{
		UnregisterAllCommands(pDTE2);
		UnregisterAllCommandBars(pDTE2);
	}
}

HRESULT CConnect::FindToolbarByName(_CommandBars* pCommandBars, LPCWSTR szToolbarName, CommandBar** pOut)
{
	if(pCommandBars == NULL || pOut == NULL)
		return E_FAIL;

	CComVariant varToolbarName(szToolbarName);
	int nToolBars = 0;
	HRESULT hr = pCommandBars->get_Count(&nToolBars);
	for(int idx = 1; idx <= nToolBars; idx++)
	{
		CComPtr<CommandBar> pCommandBar;
		hr = pCommandBars->get_Item(CComVariant(idx), &pCommandBar);
		if(FAILED(hr))
			continue;

		BSTR bsName = NULL;
		if(pCommandBar != NULL)
			pCommandBar->get_Name(&bsName);
		if(CComVariant(bsName) == varToolbarName)
		{
			*pOut = pCommandBar;
			(*pOut)->AddRef();
			break;
		}
		SysFreeString(bsName);
		hr = E_FAIL;
	}
	return hr;
}

// CConnect
STDMETHODIMP CConnect::OnConnection(IDispatch *pApplication, ext_ConnectMode ConnectMode, IDispatch *pAddInInst, SAFEARRAY ** /*custom*/ )
{
	HRESULT hr = S_OK;
	pApplication->QueryInterface(__uuidof(DTE2), (LPVOID*)&m_pDTE);
	pAddInInst->QueryInterface(__uuidof(EnvDTE::AddIn), (LPVOID*)&m_pAddInInstance);
	SetSink();
	
	if(ConnectMode == AddInDesignerObjects::ext_cm_CommandLine)
	{
		myMessageBox(_T("CConnect::OnConnection, CommandLine Mode"));
		return S_OK;
	}

	CComPtr<IDispatch> pDisp;
	CComPtr<EnvDTE::Commands> pCommands;
	CComPtr<Commands2> pCommands2;
	CComQIPtr<_CommandBars> pCommandBars;
	CComPtr<CommandBar> pCommandBar;

	CComPtr<EnvDTE::Windows> spWindows;
	CComPtr<EnvDTE::Window> spWindow;
	CComQIPtr<EnvDTE::OutputWindow> spOutputWindow;
	CComPtr<EnvDTE::OutputWindowPanes> spOutputWindowPanes;

	IfFailGoCheck(m_pDTE->get_Commands(&pCommands), pCommands);
	pCommands2 = pCommands;

	// Get the set of command bars for the application.
	IfFailGoCheck(m_pDTE->get_CommandBars(&pDisp), pDisp);
	pCommandBars = pDisp;

	// See if the toolbar has been created.
	bool forceRecreateOfToolbarAndCommands = false;
	hr = FindToolbarByName(pCommandBars, TOOLBAR_NAME, &pCommandBar);
	if (SUCCEEDED(hr))
	{
		CComPtr<CommandBarControls> pCommandBarControls;
		pCommandBar->get_Controls(&pCommandBarControls);

		int count;
		pCommandBarControls->get_Count(&count);

		if (count == 0)
		{
			if(myMessageBox(IDS_MISSINGTOOLBARCONTENTS, MB_YESNO) == IDYES)
			{
				forceRecreateOfToolbarAndCommands = true;
			}
		}

		pCommandBar = NULL;
	}
	else
	{
		forceRecreateOfToolbarAndCommands = true;
	}

    if(ConnectMode == 5 || forceRecreateOfToolbarAndCommands) //5 == ext_cm_UISetup
    {
		// See if the CvsIn toolbar has been created.
		hr = FindToolbarByName(pCommandBars, TOOLBAR_NAME, &pCommandBar);
		if(FAILED(hr))
		{
			pDisp = NULL;
			// The toolbar hasn't been created yet.  Add it.
			hr = pCommands->AddCommandBar(CComBSTR(TOOLBAR_NAME),
										  EnvDTE::vsCommandBarTypeToolbar,
										  NULL,
										  1,
										  &pDisp);

			// Yes, this code is unnecessary, but it serves to prove that
			// the command bar creation actually worked.
            hr = FindToolbarByName(pCommandBars, TOOLBAR_NAME, &pCommandBar);
		}

        int curToolbarPosition = 1;
		int nMaxCommand = sizeof(s_commandList)/sizeof(s_commandList[0]);
        for(int curCommand = 0; curCommand < nMaxCommand; ++curCommand)
        {
            CComPtr<EnvDTE::Command> pCreatedCommand;
            const CommandInfo* commandInfo = &s_commandList[curCommand];
            UnregisterCommand(pCommands, commandInfo->m_name);

			wstring regStr = myLoadString(commandInfo->m_cmdID);
			wstring buttonText;
            wstring toolTip;
            myExtractSubString(buttonText, regStr.c_str(), 1, '\n');
            myExtractSubString(toolTip, regStr.c_str(), 2, '\n');

			HRESULT hr2 = pCommands2->AddNamedCommand2(m_pAddInInstance,
										CComBSTR(commandInfo->m_name),
										CComBSTR(buttonText.c_str()),
										CComBSTR(toolTip.c_str()),
										VARIANT_FALSE,
										CComVariant(commandInfo->m_bitmapID),
										NULL,
										(EnvDTE::vsCommandStatusSupported + EnvDTE::vsCommandStatusEnabled),
										EnvDTE80::vsCommandStylePict,
										EnvDTE80::vsCommandControlTypeButton,
										&pCreatedCommand);
			if(SUCCEEDED(hr2) && (pCreatedCommand) &&  commandInfo->m_hasIcon)
			{
				//Add the control:
				pDisp = NULL;
				IfFailGoCheck(pCreatedCommand->AddControl(pCommandBar, curToolbarPosition, &pDisp),pDisp);
				curToolbarPosition++;
            }
        }
    }

	// Create a our pane from output window.
	IfFailGoCheck(m_pDTE->get_Windows(&spWindows), spWindows);
	IfFailGoCheck(spWindows->Item(CComVariant(CComBSTR(EnvDTE::vsWindowKindOutput)), &spWindow), spWindow);

	pDisp = NULL;
	IfFailGoCheck(spWindow->get_Object(&pDisp), pDisp);

	spOutputWindow = pDisp;
	if (spOutputWindow == NULL)
		return E_FAIL;

	IfFailGoCheck(spOutputWindow->get_OutputWindowPanes(&spOutputWindowPanes), spOutputWindowPanes);

	hr = spOutputWindowPanes->Item(CComVariant(OUTPANE_NAME),  &m_pOutputWindowPane);

	if (FAILED(hr))
	{
		IfFailGoCheck(spOutputWindowPanes->Add(CComBSTR(OUTPANE_NAME), &m_pOutputWindowPane), m_pOutputWindowPane);
	}

	// Print the welcome message
	PrintToOutput(L"Welcome to use CppClean AddIn for VS2008!");

	TCHAR szModulePath[_MAX_PATH] = {0};
	GetModuleFileName(_AtlModule.GetResourceInstance(), szModulePath, _MAX_PATH);
	PrintToOutput(szModulePath);

Error:

	return hr;
}

STDMETHODIMP CConnect::OnDisconnection(ext_DisconnectMode /*RemoveMode*/, SAFEARRAY ** /*custom*/ )
{
	RemoveSink();
	m_pDTE = NULL;
	m_pAddInInstance = NULL;

	return S_OK;
}

STDMETHODIMP CConnect::OnAddInsUpdate (SAFEARRAY ** /*custom*/ )
{
	return S_OK;
}

STDMETHODIMP CConnect::OnStartupComplete (SAFEARRAY ** /*custom*/ )
{
	return S_OK;
}

STDMETHODIMP CConnect::OnBeginShutdown (SAFEARRAY ** /*custom*/ )
{
	return S_OK;
}


STDMETHODIMP CConnect::QueryStatus(BSTR bstrCmdName,
								   EnvDTE::vsCommandStatusTextWanted NeededText,
								   EnvDTE::vsCommandStatus * pStatusOption,
								   VARIANT * pvarCommandText)
{
	if(NeededText == EnvDTE::vsCommandStatusTextWantedNone)
	{
		int nMaxCommand = sizeof(s_commandList)/sizeof(s_commandList[0]);
		for (int curCommand = 0; curCommand < nMaxCommand; ++curCommand)
		{
			const CommandInfo* commandInfo = &s_commandList[curCommand];
			WCHAR item[_MAX_PATH] = {0};
			wcscpy_s(item, MAX_PATH, MODULE_NAME);
			wcscat_s(item, MAX_PATH, commandInfo->m_name);
			if(!_wcsicmp(bstrCmdName, item))
			{
				*pStatusOption =(EnvDTE::vsCommandStatus)(EnvDTE::vsCommandStatusEnabled + EnvDTE::vsCommandStatusSupported);
			}
		}
	}
	return S_OK;
}

STDMETHODIMP CConnect::Exec(BSTR bstrCmdName, 
							EnvDTE::vsCommandExecOption ExecuteOption, 
							VARIANT * /*pvarVariantIn*/, 
							VARIANT * /*pvarVariantOut*/, 
							VARIANT_BOOL *pvbHandled)
{
	int curCommand;

	*pvbHandled = VARIANT_FALSE;
	if(ExecuteOption == EnvDTE::vsCommandExecOptionDoDefault)
	{
		WCHAR item[_MAX_PATH] = {0};
		int nMaxCommand = sizeof(s_commandList)/sizeof(s_commandList[0]);
		for (curCommand = 0; curCommand < nMaxCommand; ++curCommand)
		{
			const CommandInfo* commandInfo = &s_commandList[curCommand];
			wcscpy_s(item, MAX_PATH, MODULE_NAME);
			wcscat_s(item, MAX_PATH, commandInfo->m_name);
			if(!_wcsicmp(bstrCmdName, item))
				break;
		}

		if (curCommand == nMaxCommand)
		{
			return S_OK;
		}

		m_pTestSuite->HandleCommand(curCommand);

		*pvbHandled = VARIANT_TRUE;
	}
	return S_OK;
}

void CConnect::PrintToOutput(LPCTSTR format, BOOL bNewLine, ...)
{
	va_list argptr;
	va_start(argptr, bNewLine);
	TCHAR szData[MAX_STRING_LEN] = {0};
	_vstprintf_s(szData, MAX_STRING_LEN - 1, format, argptr);
	va_end(argptr);
	if(bNewLine)
		_tcscat_s(szData, MAX_STRING_LEN, _T("\r\n"));

	PrintToOutputMessage(szData);
}

/*!
Print the message to the output window
\param lpszMessage Message to be printed
\param bPrefix Set to TRUE to display the CvsIn prefix, FALSE to skip prefix
*/
void CConnect::PrintToOutputMessage(LPCTSTR lpszMessage)
{
	if (m_pOutputWindowPane)
	{
		m_pOutputWindowPane->Activate();
		m_pOutputWindowPane->OutputString(CComBSTR(lpszMessage));
	}
}


//CDebugEventsSink
HRESULT CConnect::SetSink()
{
	HRESULT hr = S_OK;
	// Get DTE.Events interface
	CComPtr<EnvDTE::Events> pEvents;
	IfFailGo(m_pDTE->get_Events(&pEvents));

	// Get DTE.BuildEvents interface
	IfFailGo(pEvents->get_DebuggerEvents((EnvDTE::_DebuggerEvents**)&m_pDebugEvents));
	IfFailGo(pEvents->get_BuildEvents((EnvDTE::_BuildEvents**)&m_pBuildEvents));

	// Sink the BuildEvents
	IfFailGo(DebuggerEventImpl::DispEventAdvise((IUnknown*)m_pDebugEvents.p));
	IfFailGo(BuildEventImpl::DispEventAdvise((IUnknown*)m_pBuildEvents.p));

Error:
	return hr;
}

void CConnect::RemoveSink()
{
	// Remove the BuildEvents Sink
	if(m_pDebugEvents)
		DebuggerEventImpl::DispEventUnadvise((IUnknown*)m_pDebugEvents.p);

	if(m_pBuildEvents)
		BuildEventImpl::DispEventUnadvise((IUnknown*)m_pBuildEvents.p);
}

void CConnect::OnEnterRunMode(EnvDTE::dbgEventReason Reason)
{
}

void CConnect::OnEnterDesignMode(EnvDTE::dbgEventReason Reason)
{
//	if(Reason == dbgEventReasonEndProgram && m_pTestSuite != NULL)
//		m_pTestSuite->StopCoverage();
}

void CConnect::OnEnterBreakMode(EnvDTE::dbgEventReason Reason, EnvDTE::dbgExecutionAction* ExecutionAction)
{

}

void CConnect::OnExceptionThrown(BSTR ExceptionType, 
								 BSTR Name, 
								 long Code, 
								 BSTR Description, 
								 EnvDTE::dbgExceptionAction* ExceptionAction)
{
}

void CConnect::OnExceptionNotHandled(BSTR ExceptionType, 
									 BSTR Name, 
									 long Code, 
									 BSTR Description, 
									 EnvDTE::dbgExceptionAction* ExceptionAction)
{
}

void CConnect::OnContextChanged(EnvDTE::Process* NewProcess, 
								EnvDTE::Program* NewProgram, 
								EnvDTE::Thread* NewThread, 
								EnvDTE::StackFrame* NewStackFrame)
{
}

void CConnect::OnBuildBegin(EnvDTE::vsBuildScope Scope, EnvDTE::vsBuildAction Action)
{

}

void CConnect::OnBuildDone(EnvDTE::vsBuildScope Scope, EnvDTE::vsBuildAction Action)
{
	if(Action == EnvDTE::vsBuildActionBuild && m_pTestSuite != NULL)
		m_pTestSuite->BuildDone();
}

void CConnect::OnBuildProjConfigBegin(BSTR Project, BSTR ProjectConfig, BSTR Platform, BSTR SolutionConfig)
{

}

void CConnect::OnBuildProjConfigDone(BSTR Project, BSTR ProjectConfig, BSTR Platform, BSTR SolutionConfig, VARIANT_BOOL Success)
{

}

